import cv2
import numpy as np

# Constants
RE_WGS84 = 6378137.0  # Earth semimajor axis (WGS84) (m)
FE_WGS84 = 1.0 / 298.257223563  # Earth flattening (WGS84)
PI = np.pi  # pi
D2R = PI / 180.0  # Degrees to radians
R2D = 180.0 / PI  # Radians to degrees


# Function to transform geodetic to ECEF position
def pos2ecef(pos):
    lat, lon, h = pos
    sin_lat = np.sin(lat)
    cos_lat = np.cos(lat)
    sin_lon = np.sin(lon)
    cos_lon = np.cos(lon)

    e2 = FE_WGS84 * (2.0 - FE_WGS84)
    N = RE_WGS84 / np.sqrt(1.0 - e2 * sin_lat ** 2)

    x = (N + h) * cos_lat * cos_lon
    y = (N + h) * cos_lat * sin_lon
    z = ((1 - e2) * N + h) * sin_lat

    return np.array([x, y, z])


# Function to transform ECEF to local tangential coordinates (ENU)
def ecef2enu(pos, vec):
    lat, lon = pos[:2]
    sin_lat = np.sin(lat)
    cos_lat = np.cos(lat)
    sin_lon = np.sin(lon)
    cos_lon = np.cos(lon)

    E = np.array([[-sin_lon, cos_lon, 0],
                  [-sin_lat * cos_lon, -sin_lat * sin_lon, cos_lat],
                  [cos_lat * cos_lon, cos_lat * sin_lon, sin_lat]])

    return E @ vec


# Function to transform LLA to ENU coordinates
def tfb_lla2enu(ls0, ls):
    pos0 = np.array([ls0[0] * D2R, ls0[1] * D2R, ls0[2]])
    pos = np.array([ls[0] * D2R, ls[1] * D2R, ls[2]])

    r0 = pos2ecef(pos0)
    r = pos2ecef(pos)
    r_diff = r - r0

    e = ecef2enu(pos0, r_diff)

    return e


def is_valid_coordinate(lat, lng):
    """
    Check if the given latitude and longitude are valid.
    Latitude is valid if it's between -85.0511 and 85.0511 degrees.
    Longitude is valid if it's between -180 and 180 degrees.
    """
    return -85.0511 <= lat <= 85.0511 and -180 <= lng <= 180


def latlng_to_tile(lat, lng, zoom):
    """
    Converts latitude and longitude to tile coordinates at the given zoom level.
    """
    n = 2.0 ** zoom
    x_tile = int((lng + 180.0) / 360.0 * n)
    y_tile = int((1.0 - np.log(np.tan(np.radians(lat)) + 1 / np.cos(np.radians(lat))) / np.pi) / 2.0 * n)
    return x_tile, y_tile


def tile_to_latlng(x_tile, y_tile, zoom):
    """
    Converts tile coordinates (x_tile, y_tile) at a given zoom level to
    latitude and longitude of the tile's upper-left corner.
    """
    n = 2.0 ** zoom
    lon_deg = x_tile / n * 360.0 - 180.0
    lat_rad = np.arctan(np.sinh(np.pi * (1 - 2 * y_tile / n)))
    lat_deg = np.degrees(lat_rad)
    return lat_deg, lon_deg


def stitch_tiles(path, topleft_lat, topleft_lng, bottomright_lat, bottomright_lng, zoom, fill_color=(255, 255, 255),
                 border_color=(0, 0, 0), border_thickness=0):
    """
    Stitch tiles from top left to bottom right coordinates.
    """
    if not (is_valid_coordinate(topleft_lat, topleft_lng) and is_valid_coordinate(bottomright_lat, bottomright_lng)):
        raise ValueError("Invalid latitude or longitude values.")

    if not 0 <= zoom <= 19:  # Assuming zoom level range is 0 to 19
        raise ValueError("Invalid zoom level.")

    # Calculate tile coordinates
    top_left_x, top_left_y = latlng_to_tile(topleft_lat, topleft_lng, zoom)
    bottom_right_x, bottom_right_y = latlng_to_tile(bottomright_lat, bottomright_lng, zoom)

    # Handle wrap around scenario
    if bottomright_lng < topleft_lng:
        bottom_right_x += 2 ** zoom

    stitched_image = None
    tile_size = (256, 256)  # Assuming each tile size is 256x256 pixels

    for y in range(top_left_y, bottom_right_y + 1):
        row_images = []
        for x in range(top_left_x, bottom_right_x + 1):
            tile_path = path + f"\\{zoom}\\{y}\\{x}.png"
            tile = cv2.imread(tile_path)
            if tile is None:
                print(f"Tile missing at {x}, {y}. Filling with color {fill_color}.")
                tile = np.full((tile_size[1], tile_size[0], 3), fill_color, dtype=np.uint8)
            elif border_thickness > 0:
                # Draw a border around the tile
                cv2.rectangle(tile, (0, 0), (tile_size[0] - 1, tile_size[1] - 1), border_color, border_thickness)
            row_images.append(tile)

        # Concatenate all images in the row
        row_image = cv2.hconcat(row_images)
        stitched_image = row_image if stitched_image is None else cv2.vconcat([stitched_image, row_image])

    # 计算左上角和右下角的经纬度
    topleft_lat, topleft_lng = tile_to_latlng(top_left_x, top_left_y, zoom)
    bottomright_lat, bottomright_lng = tile_to_latlng(bottom_right_x + 1, bottom_right_y + 1, zoom)

    return stitched_image, (topleft_lat, topleft_lng, bottomright_lat, bottomright_lng)


def binarization_by_color(input_image, color_list, threshold):
    # 将输入图像转换为float32以进行向量化运算
    input_float = np.float32(input_image)
    output = np.full(input_image.shape[:2], 255, dtype=np.uint8)  # 初始化全白图像

    # 将颜色列表转换为NumPy数组
    colors = np.array(color_list, dtype=np.float32)

    # 为每种颜色计算距离并应用阈值
    for color in colors:
        # 计算距离
        distance = np.sum(np.abs(input_float - color), axis=2)

        # 应用阈值
        mask = distance < threshold
        output[mask] = 0

    return output


def overlay_binary_on_original(original, binary, alpha=0.4):
    """
    Overlay the binary image on the original image with a given alpha transparency.
    """
    # 将二值化图像转换为与原始图像相同的三通道格式
    binary_colored = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)

    # 将二值化图像以特定透明度叠加在原始图像上
    overlay = cv2.addWeighted(original, 1 - alpha, binary_colored, alpha, 0)

    return overlay


def nautical_chart_to_grid(chat_path, topleft_lat, topleft_lng, bottomright_lat, bottomright_lng, zoom_level,
                           allow_color_list, binarization_threshold, pix_size):
    """
    将海图转换为栅格地图
    """
    stitched_image, lt_rb = stitch_tiles(chat_path, topleft_lat, topleft_lng, bottomright_lat, bottomright_lng,
                                         zoom_level, border_thickness=0)
    binary_image = binarization_by_color(stitched_image, allow_color_list, binarization_threshold)

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (4, 4))
    opened_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel, iterations=6)
    dilated_image = cv2.dilate(opened_image, kernel, iterations=4)

    enu_rb = tfb_lla2enu((lt_rb[0], lt_rb[1], 0), (lt_rb[2], lt_rb[3], 0))
    width = abs(enu_rb[0])
    heigh = abs(enu_rb[1])
    new_img_width = int(width / pix_size)
    new_img_height = int(heigh / pix_size)
    # 将图像缩放为指定大小
    resized_image = cv2.resize(dilated_image, (new_img_width, new_img_height), interpolation=cv2.INTER_AREA)
    # 将resized_image转换为栅格地图，使用0表示通行，1表示障碍物
    _, binary_image = cv2.threshold(resized_image, 1, 1, cv2.THRESH_BINARY)
    # 将图像进行翻转
    flipped_image = cv2.flip(binary_image, 0)

    return flipped_image.T
